import chisel3._
import chisel3.util._

class MyDCache(val SIZE: Int, val LINE_SIZE: Int, val WAY_NUM: Int, val UseLRU: Boolean)
    extends Module {
  // ============ 常用Width定义 ===========
  val INDEX_WIDTH   = log2Ceil(SIZE)
  val OFFSET_WIDTH  = log2Ceil(LINE_SIZE)
  val WAY_NUM_WIDTH = log2Ceil(WAY_NUM)
  val LINE_BIT_SIZE = LINE_SIZE * 8
  val LINE_WORD_NUM = LINE_SIZE / 4
  println("Data Cache Size: " + SIZE)
  println("Data Cache Line Size   : " + LINE_WORD_NUM + " Word | " + LINE_SIZE + " Byte | " + LINE_BIT_SIZE + " Bit")
  println("Data Cache Way Number  : " + WAY_NUM)
  assert(
    (INDEX_WIDTH + OFFSET_WIDTH <= 12).asBool,
    "Index windth + offset width should be less than or equal 12, but given is %d\n",
    (INDEX_WIDTH + OFFSET_WIDTH).asUInt
  )

  // ============ 外部接口(io)定义 ===========
  val io = IO(new Bundle {
    // cpu端接口(双发请求)
    val valid     = Input(Vec(2, Bool()))
    val op        = Input(Vec(2, Bool()))
    val tag       = Input(Vec(2, UInt(20.W)))
    val index     = Input(Vec(2, UInt(INDEX_WIDTH.W)))
    val offset    = Input(Vec(2, UInt(OFFSET_WIDTH.W)))
    val size      = Input(Vec(2, UInt(2.W)))
    val wstrb     = Input(Vec(2, UInt(4.W)))
    val wdata     = Input(Vec(2, UInt(32.W)))
    val addr_ok   = Output(Vec(2, Bool()))
    val data_ok   = Output(Vec(2, Bool()))
    val rdata     = Output(Vec(2, UInt(32.W)))
    val uncached  = Input(Vec(2, Bool()))

    val cacop_en  = Input(Vec(2, Bool()))
    val cacop_op  = Input(Vec(2, UInt(2.W)))
    val tlb_excp_cancel_req = Input(Bool())

    // AXI接口
    val rd_req    = Output(Bool())
    val rd_type   = Output(UInt(3.W))
    val rd_addr   = Output(UInt(32.W))
    val rd_rdy    = Input(Bool())
    val ret_valid = Input(Bool())
    val ret_last  = Input(Bool())
    val ret_data  = Input(UInt(32.W))
    val wr_req    = Output(Bool())
    val wr_type   = Output(UInt(3.W))
    val wr_addr   = Output(UInt(32.W))
    val wr_wstrb  = Output(UInt(4.W))
    val wr_data   = Output(UInt(LINE_BIT_SIZE.W))
    val wr_rdy    = Input(Bool())
  })

  // ========== 计数器统计 ===========
  val total_req_counter   = RegInit(0.U(32.W))
  val total_miss_counter  = RegInit(0.U(32.W))
  val dcache_counter      = RegInit(0.U(32.W))
  dcache_counter := dcache_counter + 1.U

  // ========== 标志与存储结构 ===========
  val Dirty = RegInit(VecInit.fill(SIZE, WAY_NUM)(false.B))
  val DataBank  = Array.fill(WAY_NUM)(Module(new DataRAM(SIZE, LINE_BIT_SIZE)).io)
  val TagvBank  = Array.fill(WAY_NUM)(Module(new TagvRAM(SIZE, 20)).io)
  val Valid     = RegInit(VecInit(Seq.fill(SIZE)(0.U(WAY_NUM.W))))

  // ========== 命中判断 ===========
  val way_hit         = Wire(Vec(2, Vec(WAY_NUM, Bool())))
  val cache_hit       = Wire(Vec(2, Bool()))
  val cache_hit_way   = Wire(Vec(2, UInt(WAY_NUM_WIDTH.W)))

  // ========== 替换路相关 ==========
  val replaced_way_buffer = RegInit(VecInit(Seq.fill(2)(0.U(WAY_NUM_WIDTH.W))))
  val replaced_index_buffer = RegInit(VecInit(Seq.fill(2)(0.U(INDEX_WIDTH.W))))
  val replaced_way = Wire(Vec(2, UInt(WAY_NUM_WIDTH.W)))
  val replaced_data = Wire(Vec(2, UInt(LINE_BIT_SIZE.W)))
  val replaced_tag = Wire(Vec(2, UInt(20.W)))
  val replaced_index = Wire(Vec(2, UInt(INDEX_WIDTH.W)))
  val replaced_addr = Wire(Vec(2, UInt(32.W)))
  val replaced_dirty = Wire(Vec(2, Bool()))
  val replaced_valid = Wire(Vec(2, Bool()))

  val need_write_axi = RegInit(VecInit(Seq.fill(2)(false.B)))
  val need_read_axi  = RegInit(VecInit(Seq.fill(2)(false.B)))

  // ========== 读/写数据通道 ===========
  val read_data = Wire(Vec(WAY_NUM, UInt(LINE_BIT_SIZE.W)))
  val target_data0 = Wire(Vec(LINE_WORD_NUM, UInt(32.W)))
  val target_data1 = Wire(Vec(LINE_WORD_NUM, UInt(32.W)))
  val read_tagv = Wire(Vec(WAY_NUM, UInt(21.W)))
  val read_tag = Wire(Vec(WAY_NUM, UInt(20.W)))
  val read_valid = Wire(Vec(WAY_NUM, Bool()))

  // ========== 请求状态/缓冲 ==========
  val receive_request = Wire(Vec(2, Bool()))
  val uncached_write = Wire(Vec(2, Bool()))
  val uncached_read = Wire(Vec(2, Bool()))
  val uncached_en = Wire(Vec(2, Bool()))
  val read_hit = Wire(Vec(2, Bool()))
  val write_hit = Wire(Vec(2, Bool()))
  val cache_unbussy = Wire(Vec(2, Bool()))

  // ========== 双发状态机 ============
  object CacheState extends ChiselEnum {
    val IDLE, LOOKUP, WRITEHIT, MISS, REPLACE, REFILL, RESPOND = Value
  }
  val main_state = RegInit(VecInit(Seq.fill(2)(CacheState.IDLE))) // 每个请求独立状态机

  // 状态机便捷信号
  val is_idle     = main_state.map(_ === CacheState.IDLE)
  val is_lookup   = main_state.map(_ === CacheState.LOOKUP)
  val is_writehit = main_state.map(_ === CacheState.WRITEHIT)
  val is_miss     = main_state.map(_ === CacheState.MISS)
  val is_replace  = main_state.map(_ === CacheState.REPLACE)
  val is_refill   = main_state.map(_ === CacheState.REFILL)
  val is_respond  = main_state.map(_ === CacheState.RESPOND)

  // ========== AXI请求缓冲 ===========
  val rd_req_pending = RegInit(VecInit(Seq.fill(2)(false.B)))
  val wr_req_pending = RegInit(VecInit(Seq.fill(2)(false.B)))
  val rd_type_buffer = RegInit(VecInit(Seq.fill(2)(0.U(3.W))))
  val rd_addr_buffer = RegInit(VecInit(Seq.fill(2)(0.U(32.W))))
  val wr_type_buffer = RegInit(VecInit(Seq.fill(2)(0.U(3.W))))
  val wr_addr_buffer = RegInit(VecInit(Seq.fill(2)(0.U(32.W))))
  val wr_wstrb_buffer = RegInit(VecInit(Seq.fill(2)(0.U(4.W))))
  val wr_data_buffer = RegInit(VecInit(Seq.fill(2)(0.U(LINE_BIT_SIZE.W))))

  // ========== 请求缓冲器 ==========
  val request_cacop_en_buffer = RegInit(VecInit(Seq.fill(2)(false.B)))
  val request_cacop_op_buffer = RegInit(VecInit(Seq.fill(2)(0.U(2.W))))
  val request_op_buffer = RegInit(VecInit(Seq.fill(2)(false.B)))
  val request_tag_buffer = RegInit(VecInit(Seq.fill(2)(0.U(20.W))))
  val request_index_buffer = RegInit(VecInit(Seq.fill(2)(0.U(INDEX_WIDTH.W))))
  val request_offset_buffer = RegInit(VecInit(Seq.fill(2)(0.U(OFFSET_WIDTH.W))))
  val request_uncached_buffer = RegInit(VecInit(Seq.fill(2)(false.B)))
  val target_word_id = RegInit(VecInit(Seq.fill(2)(0.U((OFFSET_WIDTH-2).W))))
  val request_wdata_buffer = RegInit(VecInit(Seq.fill(2)(0.U(32.W))))
  val request_wstrb_buffer = RegInit(VecInit(Seq.fill(2)(0.U(4.W))))
  val request_size_buffer = RegInit(VecInit(Seq.fill(2)(0.U(2.W))))

  val request_addr = Wire(Vec(2, UInt(32.W)))
  val request_line_addr = Wire(Vec(2, UInt(32.W)))
  val normal_request = Wire(Vec(2, Bool()))
  for (i <- 0 until 2) {
    request_addr(i) := Cat(Mux(is_lookup(i), io.tag(i), request_tag_buffer(i)), 
                          request_index_buffer(i), request_offset_buffer(i))
    request_line_addr(i) := Cat(Mux(is_lookup(i), io.tag(i), request_tag_buffer(i)), 
                              request_index_buffer(i), 0.U(OFFSET_WIDTH.W))
  }

  // miss数据缓冲
  val ret_data_buffer = Reg(Vec(2, Vec(LINE_WORD_NUM, UInt(32.W))))
  val miss_data_ret_cnt = RegInit(VecInit(Seq.fill(2)(0.U((OFFSET_WIDTH - 2).W))))

  // ========== cacop指令处理 ==========
  val cacop_op_0 = Wire(Vec(2, Bool()))
  val cacop_op_1 = Wire(Vec(2, Bool()))
  val cacop_op_2 = Wire(Vec(2, Bool()))
  val cacop_way  = Wire(Vec(2, UInt(WAY_NUM_WIDTH.W)))
  for (i <- 0 until 2) {
    cacop_op_0(i) := request_cacop_en_buffer(i) && request_cacop_op_buffer(i) === 0.U
    cacop_op_1(i) := request_cacop_en_buffer(i) && request_cacop_op_buffer(i) === 1.U
    cacop_op_2(i) := request_cacop_en_buffer(i) && request_cacop_op_buffer(i) === 2.U
    cacop_way(i)  := request_addr(i)(WAY_NUM_WIDTH-1, 0)
  }

  // 替换路选择（轮转策略）
  for (i <- 0 until 2) {
    replaced_way(i) := Mux(cacop_op_2(i) && cache_hit(i), cache_hit_way(i), 
                         Mux(cacop_op_0(i) || cacop_op_1(i), cacop_way(i), 
                             dcache_counter(WAY_NUM_WIDTH-1, 0) + i.U)) // 双发使用不同偏移避免冲突
    replaced_data(i) := read_data(replaced_way(i))
    replaced_dirty(i) := Dirty(replaced_index(i))(replaced_way(i))
    replaced_index(i) := request_index_buffer(i)
    replaced_tag(i) := read_tag(replaced_way(i))
    replaced_valid(i) := read_valid(replaced_way(i))
    replaced_addr(i) := Cat(replaced_tag(i), replaced_index(i), 0.U(OFFSET_WIDTH.W))
  }

  // 获取Tag和Valid
  for (i <- 0 until WAY_NUM) {
    read_tag(i) := read_tagv(i)(20, 1)
    read_valid(i) := read_tagv(i)(0)
  }

  // 命中判定
  for (i <- 0 until WAY_NUM) {
    way_hit(0)(i) := (read_tag(i) === io.tag(0)) && read_valid(i)
    way_hit(1)(i) := (read_tag(i) === io.tag(1)) && read_valid(i)
  }

  for (i <- 0 until 2) {
    cache_hit(i) := way_hit(i).reduce(_||_) && !io.tlb_excp_cancel_req
    cache_hit_way(i) := OHToUInt(way_hit(i))
    uncached_en(i) := Mux(is_lookup(i), io.uncached(i), request_uncached_buffer(i)) && !request_cacop_en_buffer(i)
    uncached_read(i) := uncached_en(i) && !request_op_buffer(i)
    uncached_write(i) := uncached_en(i) && request_op_buffer(i)
    normal_request(i) := !uncached_en(i) && !request_cacop_en_buffer(i)
    read_hit(i) := cache_hit(i) && !request_op_buffer(i) && normal_request(i)
    write_hit(i) := cache_hit(i) && request_op_buffer(i) && normal_request(i)
    
    // 空闲条件：IDLE 或 LOOKUP且命中 或 RESPOND
    cache_unbussy(i) := is_idle(i) || (is_lookup(i) && cache_hit(i) && normal_request(i)) || is_respond(i)
    receive_request(i) := cache_unbussy(i) && io.valid(i)
    io.addr_ok(i) := cache_unbussy(i)
    io.data_ok(i) := (is_lookup(i) && (cache_hit(i) || request_op_buffer(i)) && !uncached_en(i)) || is_respond(i)
  }

  // 读数据输出
  target_data0 := read_data(cache_hit_way(0)).asTypeOf(target_data0)
  target_data1 := read_data(cache_hit_way(1)).asTypeOf(target_data1)
  for (i <- 0 until 2) {
    io.rdata(i) := Mux(is_respond(i), 
                      ret_data_buffer(i)(target_word_id(i)), 
                      Mux(i.U === 0.U, target_data0(target_word_id(i)), target_data1(target_word_id(i))))
  }

  // ========== AXI接口仲裁 ==========
  // 合并两个请求的AXI请求（优先级：请求0 > 请求1）
  val axi_arb = Module(new Arbiter(new AXIReq, 2))
  for (i <- 0 until 2) {
    axi_arb.io.in(i).valid := rd_req_pending(i) || wr_req_pending(i)
    axi_arb.io.in(i).bits.is_write := wr_req_pending(i)
    axi_arb.io.in(i).bits.addr := Mux(rd_req_pending(i), rd_addr_buffer(i), wr_addr_buffer(i))
    axi_arb.io.in(i).bits.size := Mux(rd_req_pending(i), rd_type_buffer(i), wr_type_buffer(i))
    axi_arb.io.in(i).bits.wstrb := wr_wstrb_buffer(i)
    axi_arb.io.in(i).bits.data := wr_data_buffer(i)
  }

  io.rd_req := axi_arb.io.out.valid && !axi_arb.io.out.bits.is_write
  io.wr_req := axi_arb.io.out.valid && axi_arb.io.out.bits.is_write
  io.rd_addr := axi_arb.io.out.bits.addr
  io.wr_addr := axi_arb.io.out.bits.addr
  io.rd_type := axi_arb.io.out.bits.size
  io.wr_type := axi_arb.io.out.bits.size
  io.wr_wstrb := axi_arb.io.out.bits.wstrb
  io.wr_data := axi_arb.io.out.bits.data

  // ========== 数据合并处理 ===========
  for (i <- 0 until 2) {
    val ret_data_v = io.ret_data.asTypeOf(Vec(4, UInt(8.W)))
    val target_data_v = Mux(i.U === 0.U, 
                           target_data0(target_word_id(0)).asTypeOf(Vec(4, UInt(8.W))),
                           target_data1(target_word_id(1)).asTypeOf(Vec(4, UInt(8.W))))
    val request_wdata_v = request_wdata_buffer(i).asTypeOf(Vec(4, UInt(8.W)))
    
    val merged_data = Wire(Vec(4, UInt(8.W)))
    for (j <- 0 until 4) {
      merged_data(j) := Mux(request_wstrb_buffer(i)(j), request_wdata_v(j), ret_data_v(j))
    }
    
    when(is_refill(i) && io.ret_valid) {
      ret_data_buffer(i)(miss_data_ret_cnt(i)) := merged_data.asUInt
    }
  }

  // ========== Bank访问控制 ===========
  for (p <- 0 until 2) {
    for (i <- 0 until WAY_NUM) {
      DataBank(i).clka := clock
      DataBank(i).clkb := clock
      
      // 写端口A：命中写或替换写
      DataBank(i).addra := Mux(write_hit(p) && way_hit(p)(i), 
                              request_index_buffer(p), 
                              replaced_index_buffer(p))
      DataBank(i).dina := Mux(is_lookup(p), 
                            Mux(p.U === 0.U, 
                                target_data0.asUInt, 
                                target_data1.asUInt), 
                            ret_data_buffer(p).asUInt)
      DataBank(i).wea := (is_lookup(p) && write_hit(p) && way_hit(p)(i)) || 
                        (is_respond(p) && normal_request(p) && (i.U === replaced_way_buffer(p)))
      
      // 读端口B：正常读
      DataBank(i).addrb := io.index(p)
      read_data(i) := DataBank(i).doutb
    }
  }

  // ========== Tag和Valid写入 ===========
  for (p <- 0 until 2) {
    for (i <- 0 until WAY_NUM) {
      TagvBank(i).clka := clock
      TagvBank(i).addra := Mux(receive_request(p), io.index(p), request_index_buffer(p))
      TagvBank(i).dina := Mux(request_cacop_en_buffer(p), 0.U(20.W), request_tag_buffer(p))
      TagvBank(i).wea := is_miss(p) && !uncached_en(p) && 
                        ((i.U === replaced_way_buffer(p)) || request_cacop_en_buffer(p))
    }
    
    // 更新Valid位
    val valid_update = Wire(Vec(WAY_NUM, Bool()))
    for (i <- 0 until WAY_NUM) {
      valid_update(i) := (is_miss(p) && !uncached_en(p) && (i.U === replaced_way_buffer(p))) || 
                        (is_respond(p) && normal_request(p))
    }
    when(valid_update.asUInt.orR) {
      Valid(request_index_buffer(p)) := Mux1H(valid_update, 
                                            (0 until WAY_NUM).map(j => (j.U === replaced_way_buffer(p)).asUInt))
    }
  }

  // ========== 双发状态机 ==========
  for (i <- 0 until 2) {
    switch(main_state(i)) {
      is(CacheState.IDLE) {
        when(receive_request(i)) {
          // 缓冲请求信息
          request_cacop_en_buffer(i) := io.cacop_en(i)
          request_cacop_op_buffer(i) := io.cacop_op(i)
          request_op_buffer(i) := io.op(i)
          request_tag_buffer(i) := io.tag(i)
          request_index_buffer(i) := io.index(i)
          request_offset_buffer(i) := io.offset(i)
          request_size_buffer(i) := io.size(i)
          request_wstrb_buffer(i) := io.wstrb(i)
          request_wdata_buffer(i) := io.wdata(i)
          target_word_id(i) := io.offset(i)(OFFSET_WIDTH-1, 2)
          main_state(i) := CacheState.LOOKUP
        }
      }
      is(CacheState.LOOKUP) {
        total_req_counter := total_req_counter + 1.U
        request_uncached_buffer(i) := io.uncached(i)
        
        when(io.tlb_excp_cancel_req) {
          main_state(i) := CacheState.IDLE
        }.elsewhen(!cache_hit(i) || request_cacop_en_buffer(i) || uncached_en(i)) {
          // Miss处理
          when(uncached_en(i)) { target_word_id(i) := 0.U }
          replaced_way_buffer(i) := replaced_way(i)
          replaced_index_buffer(i) := replaced_index(i)
          need_write_axi(i) := (cacop_op_2(i) && cache_hit(i) || cacop_op_1(i) || cacop_op_0(i) || normal_request(i)) && 
                              replaced_dirty(i) && replaced_valid(i) || uncached_write(i)
          
          // 缓冲AXI请求信息
          rd_type_buffer(i) := Mux(uncached_en(i), Cat(0.U(1.W), request_size_buffer(i)), "b100".U)
          wr_type_buffer(i) := Mux(uncached_en(i), Cat(0.U(1.W), request_size_buffer(i)), "b100".U)
          wr_wstrb_buffer(i) := Mux(uncached_en(i), request_wstrb_buffer(i), "b1111".U)
          rd_addr_buffer(i) := Mux(uncached_en(i), request_addr(i), request_line_addr(i))
          wr_addr_buffer(i) := Mux(uncached_en(i), request_addr(i), replaced_addr(i))
          wr_data_buffer(i) := Mux(uncached_en(i), request_wdata_buffer(i), replaced_data(i))
          
          main_state(i) := Mux(cacop_op_2(i) && !cache_hit(i), CacheState.IDLE, CacheState.MISS)
          total_miss_counter := total_miss_counter + 1.U
        }.elsewhen(receive_request(i)) {
          // 连续命中时提前接收新请求
          request_cacop_en_buffer(i) := io.cacop_en(i)
          request_cacop_op_buffer(i) := io.cacop_op(i)
          request_op_buffer(i) := io.op(i)
          request_tag_buffer(i) := io.tag(i)
          request_index_buffer(i) := io.index(i)
          request_offset_buffer(i) := io.offset(i)
          request_size_buffer(i) := io.size(i)
          request_wstrb_buffer(i) := io.wstrb(i)
          request_wdata_buffer(i) := io.wdata(i)
          target_word_id(i) := io.offset(i)(OFFSET_WIDTH-1, 2)
          main_state(i) := CacheState.LOOKUP
        }.otherwise {
          main_state(i) := CacheState.IDLE
        }
        
        when(write_hit(i)) {
          Dirty(request_index_buffer(i))(cache_hit_way(i)) := true.B
        }
      }
      is(CacheState.MISS) {
        when(io.wr_rdy || !need_write_axi(i)) {
          need_read_axi(i) := uncached_read(i) || normal_request(i)
          main_state(i) := CacheState.REPLACE
        }
      }
      is(CacheState.REPLACE) {
        when(need_read_axi(i)) {
          when(io.rd_rdy) {
            main_state(i) := CacheState.REFILL
            miss_data_ret_cnt(i) := 0.U
            ret_data_buffer(i) := VecInit(Seq.fill(LINE_WORD_NUM)(0.U(32.W)))
          }
        }.otherwise {
          // 无读请求时直接处理下一个请求
          when(receive_request(i)) {
            request_cacop_en_buffer(i) := io.cacop_en(i)
            request_cacop_op_buffer(i) := io.cacop_op(i)
            request_op_buffer(i) := io.op(i)
            request_tag_buffer(i) := io.tag(i)
            request_index_buffer(i) := io.index(i)
            request_offset_buffer(i) := io.offset(i)
            request_size_buffer(i) := io.size(i)
            request_wstrb_buffer(i) := io.wstrb(i)
            request_wdata_buffer(i) := io.wdata(i)
            target_word_id(i) := io.offset(i)(OFFSET_WIDTH-1, 2)
            main_state(i) := CacheState.LOOKUP
          }.otherwise {
            main_state(i) := CacheState.IDLE
          }
        }
      }
      is(CacheState.REFILL) {
        when(io.ret_valid) {
          miss_data_ret_cnt(i) := miss_data_ret_cnt(i) + 1.U
          when(io.ret_last) {
            when(normal_request(i)) {
              Dirty(replaced_index_buffer(i))(replaced_way_buffer(i)) := request_op_buffer(i)
            }
            main_state(i) := CacheState.RESPOND
          }
        }
      }
      is(CacheState.RESPOND) {
        when(receive_request(i)) {
          request_cacop_en_buffer(i) := io.cacop_en(i)
          request_cacop_op_buffer(i) := io.cacop_op(i)
          request_op_buffer(i) := io.op(i)
          request_tag_buffer(i) := io.tag(i)
          request_index_buffer(i) := io.index(i)
          request_offset_buffer(i) := io.offset(i)
          request_size_buffer(i) := io.size(i)
          request_wstrb_buffer(i) := io.wstrb(i)
          request_wdata_buffer(i) := io.wdata(i)
          target_word_id(i) := io.offset(i)(OFFSET_WIDTH-1, 2)
          main_state(i) := CacheState.LOOKUP
        }.otherwise {
          main_state(i) := CacheState.IDLE
        }
      }
    }
  }
}

// AXI请求类型定义（用于仲裁器）
class AXIReq extends Bundle {
  val is_write = Bool()
  val addr = UInt(32.W)
  val size = UInt(3.W)
  val wstrb = UInt(4.W)
  val data = UInt(256.W) // 假设LINE_BIT_SIZE=256
}

// 简单的仲裁器实现
class Arbiter(typ: AXIReq, n: Int) extends Module {
  val io = IO(new Bundle {
    val in = Flipped(Vec(n, Decoupled(typ)))
    val out = Decoupled(typ)
  })
  
  val grant = PriorityEncoder(io.in.map(_.valid))
  io.out.valid := io.in.map(_.valid).reduce(_||_)
  io.out.bits := Mux1H(grant.asUInt, io.in.map(_.bits))
  
  for (i <- 0 until n) {
    io.in(i).ready := io.out.ready && (grant === i.U)
  }
}